Review Summary

Nice feature addition — well-scoped passthrough to Megatron for fine-grained activation offloading. A few doc corrections and suggestions below.

Performance Evidence

This PR adds a memory-reduction feature but the description doesn't include benchmark numbers. Could you share peak GPU memory (GiB) and tokens/sec with and without activation offloading for a representative config (e.g. a MoE model with ["moe_act"] offloading)? Even a single run would help users know what to expect. A reference to the upstream Megatron-LM feature guide is also welcome.

Documentation

No user-facing docs added under docs/. Consider adding a brief section or pointer to the upstream Megatron-LM activation offloading guide.

Tests (optional nit)

Zero test coverage for the new validation logic. The existing TestApplyPerformanceConfig class in tests/unit/models/megatron/test_megatron_setup.py would be a natural home for a happy-path test (valid config sets attributes) and error-path test (empty list raises ValueError).

Generated by Claude Code

Self-review completed with /review-pr-team using a full base-to-head review and an adversarial second pass.

Reviewed areas:

• NeMo-RL config/provider propagation and explicit disable semantics
• pinned Megatron Bridge/Megatron-Core validation and module compatibility
• forward-only current/reference-logprob lifecycle versus activation-offload warmup/cache state
• exception restoration and VPP shared/distinct-config behavior
• documentation, exemplar/reference configs, unit coverage, and GB200 evidence

No remaining runtime correctness defect was found in the live PR head 01398467224921c058a70702cb4a8285eb98fc71. All six outstanding review threads were re-verified against that exact head, replied to with their fixing commit, and resolved. Terry's NUMA concern is covered by bind_to_gpu_numa(local_rank) from 9e01af64b3891e5bcc01885e10e9ca185b3e3690, which applies CPU affinity and NUMA memory binding to every Megatron worker.

The self-review found two documentation/config gaps: clarify that NeMo-RL currently supports the dense offload pair through the Transformer Engine CUDA-graph path even though pinned MCore also permits full_iteration, and add the disabled defaults to the standalone RM exemplar/reference config. Both are fixed in signed fork commit f651ac00ffc8a4cbe83a6abe626d6bea3f788325; Lyris job 2680213 validated its RM reference-config contract on Linux (1 passed, 13 deselected). The current PR branch is hosted in NVIDIA-NeMo/RL, so a maintainer must integrate that fork commit into this PR head.

GB200 validation and honest claim boundaries are now documented in the PR body. The Qwen3-30B-A3B ON arm offloaded nonzero moe_act volume on all 16 ranks and completed 3-step and 10-step lifecycle gates with finite training metrics. The single matched pair did not show a speedup, and allocator peak memory was not measured, so no stable throughput or peak-memory benefit is claimed.

Team review of the activation-offloading feature — seven agents over the diff, upstream Megatron-LM/Bridge sources at the pinned SHAs, and the full PR thread history, followed by an adversarial pass that challenged every finding.

Overall: this is a well-built PR and close to ready. The _suspend_activation_offload_for_forward_only context manager is correctly shaped — module-level, fast path first, restore-on-exception right, and it matches MCore's own pre/post-warmup-hook idiom. The 300 lines of new tests need no GPU, no Ray, and no real Megatron model, which is a genuinely good seam. All eleven earlier review threads check out as fixed at head. The knob defaults to false in every exemplar, so there is no regression risk for anyone who does not opt in.

One item I would hold for: the PPO value worker's forward-only pass never enters the new context manager, and this PR ships the value.megatron_cfg keys that make that path reachable. Details inline on train.py.

The rest is one narrow None-vs-[] regression, one docs correction worth making (the CUDA-graph requirement is stated unconditionally but only applies when CUDA graphs are on), and a handful of nits you should feel free to take or leave.

On performance evidence: the GB200 validation and the explicit statement that it does not establish a speedup or a peak-memory reduction are the honest way to present this, and that is appreciated rather than criticized. The one gap worth a follow-up, whenever convenient and not as a merge condition: the cluster runs used force_on_policy_ratio=true, so the current-policy forward-only path — the reason the new context manager exists — was only exercised by mocked unit tests. A short GRPO run with force_on_policy_ratio=false, plus torch.cuda.max_memory_allocated() on and off, would both close that gap and give users a number to decide on. If a config that OOMs with the flag off and completes with it on is easy to produce, that single data point would justify the feature better than anything else.

Two notes on the record, neither a defect. The NUMA concern from the earlier thread is genuinely closed by bind_to_gpu_numa on the Megatron worker path — but the crediting commit 9e01af64b is #2613 (Youngeun Kwon), already on main; this PR touches no worker files. And the fork commit f651ac00f does not need maintainer integration: the exemplar coverage here already matches the convention set by the nearest sibling key, so its only genuinely useful half is the docs correction flagged inline, which you can push directly.

Also checked and cleared, so they do not resurface: MCore's own validation is not bypassed (setup.py:675 re-runs __post_init__() after the assignments, so bad module names and a bare attn_proj are caught upstream — no need to duplicate those checks); the new unconditional PipelineOffloadManager import is safe at the pin; model_overrides cannot silently clobber the new keys; and the num_moe_experts guard is live and adds coverage MCore genuinely lacks.

Note: repo linters could not be run in this environment (the lockfile is linux-only and this host is macOS), so there is no lint evidence in either direction.

Generated by Claude Code